"""比较和可视化不同消融实验的模型性能

此脚本用于：
1. 识别不同消融设置的训练运行
2. 评估并比较它们的性能
3. 生成对比 GIF 动画
4. 创建性能对比报告

消融组件：
- r_ext: 环境奖励（基线）
- r_nov: 新颖性奖励（novel_type, self_coef, other_coef）
- r_hin: 后见奖励（use_hdd, ir_coef, ad_coef）
"""

import argparse
import os
import json
import numpy as np
from pathlib import Path


def identify_ablation_type(config):
    """识别消融类型

    Returns:
        str: 消融类型描述
    """
    components = []

    # 检查新颖性奖励
    novel_type = config.get('novel_type', 0)
    self_coef = config.get('self_coef', 0.0)
    other_coef = config.get('other_coef', 0.0)
    has_novelty = (novel_type > 0) and (self_coef > 0 or other_coef > 0)

    # 检查后见奖励
    use_hdd = config.get('use_hdd', False)
    ir_coef = config.get('ir_coef', 0.0)
    ad_coef = config.get('ad_coef', 0.0)
    has_hindsight = use_hdd and (ir_coef > 0 or ad_coef > 0)

    # 构建描述
    if not has_novelty and not has_hindsight:
        return "baseline (r_ext only)"

    if has_novelty:
        components.append("r_nov")
    if has_hindsight:
        components.append("r_hin")

    return " + ".join(["r_ext"] + components)


def analyze_training_run(run_dir):
    """分析单个训练运行

    Returns:
        dict: 包含运行信息的字典
    """
    params_path = os.path.join(run_dir, 'params.json')

    if not os.path.exists(params_path):
        return None

    with open(params_path, 'r') as f:
        config = json.load(f)

    # 获取可用的检查点
    models_dir = os.path.join(run_dir, 'models')
    checkpoints = []
    if os.path.exists(models_dir):
        checkpoints = sorted([d for d in os.listdir(models_dir)
                            if d.startswith('cp_')])

    # 获取最终奖励
    rewards_path = os.path.join(run_dir, 'rewards.txt')
    final_reward = None
    if os.path.exists(rewards_path):
        with open(rewards_path, 'r') as f:
            lines = f.readlines()
            if lines:
                try:
                    final_reward = float(lines[-1].strip())
                except:
                    pass

    return {
        'run_dir': run_dir,
        'run_name': os.path.basename(run_dir),
        'ablation_type': identify_ablation_type(config),
        'config': config,
        'checkpoints': checkpoints,
        'final_reward': final_reward,
        'novel_type': config.get('novel_type', 0),
        'self_coef': config.get('self_coef', 0.0),
        'other_coef': config.get('other_coef', 0.0),
        'use_hdd': config.get('use_hdd', False),
        'ir_coef': config.get('ir_coef', 0.0),
        'ad_coef': config.get('ad_coef', 0.0),
    }


def find_all_runs(base_dir='logs/results', env_name='Overcooked-v0'):
    """查找所有训练运行"""
    runs = []
    search_path = os.path.join(base_dir, env_name)

    if not os.path.exists(search_path):
        print(f"警告: 找不到目录 {search_path}")
        return runs

    # 递归搜索所有包含 params.json 的目录
    for root, dirs, files in os.walk(search_path):
        if 'params.json' in files and 'models' in dirs:
            run_info = analyze_training_run(root)
            if run_info:
                runs.append(run_info)

    return runs


def group_by_ablation(runs):
    """按消融类型分组"""
    groups = {}
    for run in runs:
        ablation = run['ablation_type']
        if ablation not in groups:
            groups[ablation] = []
        groups[ablation].append(run)
    return groups


def print_summary(runs):
    """打印摘要"""
    print("\n" + "="*80)
    print("训练运行摘要")
    print("="*80)

    groups = group_by_ablation(runs)

    for ablation_type in sorted(groups.keys()):
        print(f"\n{ablation_type}:")
        print("-" * 80)

        for run in groups[ablation_type]:
            print(f"\n  运行: {run['run_name']}")
            print(f"  路径: {run['run_dir']}")

            if run['novel_type'] > 0:
                print(f"  新颖性: type={run['novel_type']}, " +
                      f"self_coef={run['self_coef']:.4f}, " +
                      f"other_coef={run['other_coef']:.4f}")

            if run['use_hdd']:
                print(f"  后见: ir_coef={run['ir_coef']:.4f}, " +
                      f"ad_coef={run['ad_coef']:.4f}")

            if run['final_reward'] is not None:
                print(f"  最终奖励: {run['final_reward']:.2f}")

            if run['checkpoints']:
                print(f"  检查点: {', '.join(run['checkpoints'][-5:])}")  # 最后5个

    print("\n" + "="*80)
    print(f"总计: {len(runs)} 个训练运行")
    print("="*80 + "\n")


def generate_comparison_commands(runs, output_file='compare_ablations_commands.sh'):
    """生成用于比较的命令脚本"""

    if not runs:
        print("没有找到训练运行")
        return

    groups = group_by_ablation(runs)

    commands = [
        "#!/bin/bash",
        "# 自动生成的消融实验比较命令",
        "",
        "# 为每个消融类型生成 GIF",
        ""
    ]

    for ablation_type, run_list in groups.items():
        commands.append(f"\n# {ablation_type}")
        commands.append("#" + "=" * 70)

        for run in run_list:
            if not run['checkpoints']:
                continue

            # 使用最后一个检查点
            last_cp = run['checkpoints'][-1]
            model_dir = os.path.join(run['run_dir'], 'models', last_cp)

            # 创建输出目录名（安全的文件名）
            safe_name = run['run_name'].replace(':', '_').replace(' ', '_')
            output_dir = f"renders_ablation/{safe_name}_{last_cp}"

            # 获取layout名称
            layout_name = run['config'].get('layout_name', 'small_coor_2')

            command = f"""
echo "评估: {ablation_type} - {run['run_name']}"
python generate_gif.py \\
  --model_dir {model_dir} \\
  --layout_name {layout_name} \\
  --n_episodes 3 \\
  --output_dir {output_dir} \\
  --fps 5 \\
  --deterministic
"""
            commands.append(command)

    # 添加详细评估命令
    commands.append("\n\n# 详细评估（包含轨迹数据）")
    commands.append("#" + "=" * 70)

    for ablation_type, run_list in groups.items():
        for run in run_list:
            if not run['checkpoints']:
                continue

            last_cp = run['checkpoints'][-1]
            model_dir = os.path.join(run['run_dir'], 'models', last_cp)
            safe_name = run['run_name'].replace(':', '_').replace(' ', '_')
            output_dir = f"eval_ablation/{safe_name}_{last_cp}"
            layout_name = run['config'].get('layout_name', 'small_coor_2')

            command = f"""
python eval_overcooked_detailed.py \\
  --model_dir {model_dir} \\
  --layout_name {layout_name} \\
  --n_eval_episodes 10 \\
  --output_dir {output_dir} \\
  --deterministic
"""
            commands.append(command)

    # 写入文件
    with open(output_file, 'w') as f:
        f.write('\n'.join(commands))

    # 设置可执行权限
    os.chmod(output_file, 0o755)

    print(f"\n✓ 已生成命令脚本: {output_file}")
    print(f"  运行: ./{output_file}")
    print(f"  或查看文件内容选择性运行特定命令\n")


def generate_comparison_report(runs, output_file='ablation_comparison_report.txt'):
    """生成详细的对比报告"""

    groups = group_by_ablation(runs)

    lines = [
        "="*80,
        "MACE 消融实验对比报告",
        "="*80,
        "",
        f"总训练运行数: {len(runs)}",
        f"消融类型数: {len(groups)}",
        "",
    ]

    # 按消融类型汇总
    lines.append("消融类型分布:")
    lines.append("-" * 80)
    for ablation_type in sorted(groups.keys()):
        count = len(groups[ablation_type])
        lines.append(f"  {ablation_type}: {count} 个运行")
    lines.append("")

    # 详细信息
    lines.append("="*80)
    lines.append("详细信息")
    lines.append("="*80)

    for ablation_type in sorted(groups.keys()):
        lines.append(f"\n{ablation_type}")
        lines.append("-" * 80)

        for i, run in enumerate(groups[ablation_type], 1):
            lines.append(f"\n运行 {i}: {run['run_name']}")
            lines.append(f"  路径: {run['run_dir']}")

            # 配置参数
            lines.append(f"\n  配置:")
            lines.append(f"    novel_type: {run['novel_type']}")
            lines.append(f"    self_coef: {run['self_coef']:.6f}")
            lines.append(f"    other_coef: {run['other_coef']:.6f}")
            lines.append(f"    use_hdd: {run['use_hdd']}")
            lines.append(f"    ir_coef: {run['ir_coef']:.6f}")
            lines.append(f"    ad_coef: {run['ad_coef']:.6f}")

            # 训练结果
            if run['final_reward'] is not None:
                lines.append(f"\n  性能:")
                lines.append(f"    最终奖励: {run['final_reward']:.2f}")

            # 可用检查点
            if run['checkpoints']:
                lines.append(f"\n  可用检查点 ({len(run['checkpoints'])}):")
                lines.append(f"    {', '.join(run['checkpoints'])}")

    # 推荐命令
    lines.append("\n\n" + "="*80)
    lines.append("推荐命令")
    lines.append("="*80)
    lines.append("\n1. 生成对比 GIF:")
    lines.append("   ./compare_ablations_commands.sh")
    lines.append("\n2. 查看特定运行的详细评估:")
    lines.append("   python eval_overcooked_detailed.py --model_dir <model_dir> --deterministic")
    lines.append("\n3. 分析轨迹数据:")
    lines.append("   python analyze_trajectories.py <trajectories.npz>")

    lines.append("\n" + "="*80)
    lines.append("报告生成完成")
    lines.append("="*80 + "\n")

    # 写入文件
    with open(output_file, 'w') as f:
        f.write('\n'.join(lines))

    print(f"✓ 已生成对比报告: {output_file}\n")


def main():
    parser = argparse.ArgumentParser(
        description='比较和分析 MACE 消融实验',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:

  # 查看所有训练运行
  python compare_ablations.py --summary

  # 生成比较命令脚本
  python compare_ablations.py --generate-commands

  # 生成完整报告
  python compare_ablations.py --generate-report

  # 全部生成
  python compare_ablations.py --all
        """
    )

    parser.add_argument('--base-dir', type=str, default='logs/results',
                       help='训练日志基础目录')
    parser.add_argument('--env-name', type=str, default='Overcooked-v0',
                       help='环境名称')
    parser.add_argument('--summary', action='store_true',
                       help='打印训练运行摘要')
    parser.add_argument('--generate-commands', action='store_true',
                       help='生成比较命令脚本')
    parser.add_argument('--generate-report', action='store_true',
                       help='生成详细对比报告')
    parser.add_argument('--all', action='store_true',
                       help='执行所有操作')

    args = parser.parse_args()

    # 如果没有指定任何操作，默认显示摘要
    if not (args.summary or args.generate_commands or args.generate_report or args.all):
        args.summary = True

    print("\n" + "="*80)
    print("MACE 消融实验比较工具")
    print("="*80)
    print(f"\n正在搜索: {args.base_dir}/{args.env_name}/")

    # 查找所有训练运行
    runs = find_all_runs(args.base_dir, args.env_name)

    if not runs:
        print("\n⚠ 未找到任何训练运行")
        print(f"请检查路径: {args.base_dir}/{args.env_name}/\n")
        return

    print(f"✓ 找到 {len(runs)} 个训练运行\n")

    # 执行请求的操作
    if args.summary or args.all:
        print_summary(runs)

    if args.generate_commands or args.all:
        generate_comparison_commands(runs)

    if args.generate_report or args.all:
        generate_comparison_report(runs)

    print("\n完成！\n")


if __name__ == "__main__":
    main()
